# coding=utf8
from __future__ import division
import torch
import os, time
from torch.autograd import Variable
import logging
import torch.nn.functional as F
from logs import configlog
from torch import nn
import torch
import torch.utils.data as torchdata
from models.spherenet import sphere20a
import numpy as np
from convert_weight.incep_res_v1_converter import tf_to_ptch
from metrics import *
from dataset.facedata import LFW_np,LFW_zip
from datetime import datetime

model = sphere20a()
model.eval()

x1 = torch.FloatTensor(10, 3, 112, 96)
x1 = torch.autograd.Variable(x1)
x2 = torch.FloatTensor(10, 3, 112, 96)
x2 = torch.autograd.Variable(x2)

loaded_weight = torch.load('/home/hszc/zhangchi/channel-prune/prune_result/sphere2/temp/temp_bestmodel.pth')
loaded_weight = torch.load('/home/hszc/zhangchi/sphereface_pytorch/model/sphere20a_20171020.pth')
model.load_state_dict(loaded_weight)
# for name1,module1 in model._modules.items():
#     if name1.startswith('conv'):
#         weight_name = ('.').join(['module',name1,'weight'])
#         bias_name = ('.').join(['module',name1,'bias'])
#
#         new_in = loaded_weight[weight_name].size(1)
#         new_out = loaded_weight[weight_name].size(0)
#
#         module1.in_channels = new_in
#         module1.out_channels = new_out
#         module1.weight.data = loaded_weight[weight_name]
#         module1.bias.data = loaded_weight[bias_name]
#     elif name1.startswith('relu'):
#         weight_name = ('.').join(['module', name1, 'weight'])
#         new_num_para = loaded_weight[weight_name].size(0)
#
#         module1.num_parameters = new_num_para
#         module1.weight.data = loaded_weight[weight_name]
#
#     elif name1 == 'fc5':
#         weight_name = ('.').join(['module',name1,'weight'])
#         bias_name = ('.').join(['module',name1,'bias'])
#
#         new_in = loaded_weight[weight_name].size(1)
#         new_out = loaded_weight[weight_name].size(0)
#
#         module1.in_features = new_in
#         module1.out_features = new_out
#         module1.weight.data = loaded_weight[weight_name]
#         module1.bias.data = loaded_weight[bias_name]
#
#     elif name1=='fc6':
#         weight_name = ('.').join(['module', name1, 'weight'])
#         module1.weight.data = loaded_weight[weight_name]
#
# torch.save(model,'../sphere_prune.pth')

model.cuda()
model.eval()
model.feature = True


batch_size = 32
usecuda = 1

dataset = {}
data_loader={}
dataset['val'] = LFW_zip( lfw_path='/media/hszc/data/lfw.zip',
                     pairtxt='/home/hszc/zhangchi/sphereface_pytorch/data/pairs.txt',
                     landmarktxt='/home/hszc/zhangchi/sphereface_pytorch/data/lfw_landmark.txt')

data_loader['val'] = torchdata.DataLoader(dataset['val'] ,batch_size=batch_size, num_workers=4,
                                          shuffle=False, pin_memory=True)


scores_all = np.zeros(len(dataset['val'] ))
labels_all =  np.zeros(len(dataset['val'] ))


epoch_size = len(dataset['val']) // batch_size

t0 = time.time()
idx=0
for batch_cnt, data in enumerate(data_loader['val']):
    print '%s/%s'%(batch_cnt,epoch_size)
    t1 = time.time()
    since = t1 - t0
    t0 = t1
    imgs1,imgs2, labels = data

    if usecuda >= 1:
        imgs1 = Variable(imgs1.cuda())
        imgs2 = Variable(imgs2.cuda())
        labels = Variable(labels.cuda())

    else:
        imgs1 = Variable(imgs1)
        imgs2 = Variable(imgs2)
        labels = Variable(labels.cuda())
    #
    # # forward
    emd1 = model(imgs1)
    emd2 = model(imgs2)

    scores = torch.sum(emd1*emd2,dim=1) / \
          (torch.norm(emd1, p=2,dim=1) * torch.norm(emd2, p=2,dim=1) + 1e-5)
    scores = scores.data.cpu().numpy()

    # emd1 = emd1.data.cpu().numpy()
    # emd2 = emd2.data.cpu().numpy()
    # labels = labels.data.cpu().numpy()
    # scores = cal_cos(emd1,emd2)

    scores_all[idx:idx+labels.shape[0]] = scores
    labels_all[idx:idx+labels.shape[0]] = labels

    idx = idx + labels.shape[0]

fpr, tpr, thresh = roc_curve(labels_all, scores_all)
auc_score = auc(fpr, tpr)
best_acc, best_thresh = find_best_acc(labels_all, scores_all)

print 'auc_score',auc_score
print 'best_acc',best_acc
print 'best_thresh',best_thresh


